package com.todo.common.utils;


import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.*;
import java.util.function.Function;

/**
 * @author azhebuxing
 * @date 2024/4/12 11:13
 * @description
 */
@Slf4j
public class BatchExecuteUtils {


    private static final int DEFAULT_SIZE = 5;
    private static final int DEFAULT_TIMEOUT = 10000;
//*****************************  单个加载数据函数  ***********************************

    /**
     * 多线程执行，最后返回结果
     *
     * @param function 实际调用方法
     * @param params   参数集合List<P>
     * @param timeout  总最大超时时间 ms
     * @param executor 线程池
     * @param <P>      参数类型
     * @param <R>      返回类型
     * @return 返回结果
     */
    public static <P, R> List<R> executeBySingle(Function<P, R> function, List<P> params, int timeout, TimeUnit unit, Executor executor) throws InterruptedException, ExecutionException, TimeoutException {
        if (CollectionUtils.isEmpty(params)) {
            return Collections.emptyList();
        }
        List<CompletableFuture<R>> futureList = Lists.newArrayListWithCapacity(params.size());
        params.forEach(item -> futureList.add(CompletableFuture.supplyAsync(() -> function.apply(item), executor)));
        return getResultByExecutors(futureList, timeout, unit);
    }

    public static <R> List<R> getResultByExecutors(List<CompletableFuture<R>> futureList, int timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        CompletableFuture<List<R>> completableFuture = CompletableFuture.allOf(futureList.toArray(new CompletableFuture[0]))
                .thenApply((Void) -> {
                    List<R> resultList = Lists.newArrayListWithCapacity(futureList.size());
                    futureList.forEach(future -> {
                        R r = future.join();
                        if (Objects.nonNull(r)){
                            resultList.add(r);
                        }
                    });
                    return resultList;
                });
        return completableFuture.get(timeout, unit);
    }


//*****************************  批量加载数据函数  ***********************************

    /**
     * 对 List<List> 数据结构进行处理，平铺成一个List返回
     *
     * @param list
     * @param <R>
     * @return
     */
    public static <R> List<R> getList(List<List<R>> list) {
        List<R> allList = new ArrayList<>();
        for (List<R> rs : list) {
            if (!CollectionUtils.isEmpty(rs)) {
                allList.addAll(rs);
            }
        }
        return allList;
    }

    /**
     * 串行同步分批次执行
     *
     * @param function
     * @param params
     * @param groupSize
     * @param <P>
     * @param <R>
     * @return
     */
    public static <P, R> List<R> syncExecute(Function<List<P>, R> function, List<P> params, int groupSize) {
        if (CollectionUtils.isEmpty(params)) {
            return Collections.emptyList();
        }
        List<R> allList = new ArrayList<>(params.size());
        for (List<P> subParams : Lists.partition(params, groupSize)) {
            R result = function.apply(subParams);
            if (Objects.nonNull(result)) {
                allList.add(result);
            }
        }
        return allList;
    }


    /**
     * 多线程分批执行，最后返回结果
     *
     * @param function  实际调用方法
     * @param params    参数集合List<P>
     * @param groupSize 每批的大小
     * @param executor  线程池
     * @param <P>       参数类型
     * @param <R>       返回类型
     * @return 返回结果
     */
    public static <P, R> List<R> execute(Function<List<P>, R> function, List<P> params, int groupSize, Executor executor) {
        return execute(function, params, groupSize, DEFAULT_TIMEOUT, executor);
    }

    /**
     * 多线程分批执行，最后返回结果
     *
     * @param function  实际调用方法
     * @param params    参数集合List<P>
     * @param groupSize 每批的大小
     * @param timeout   总最大超时时间 ms
     * @param executor  线程池
     * @param <P>       参数类型
     * @param <R>       返回类型
     * @return 返回结果
     */
    public static <P, R> List<R> execute(Function<List<P>, R> function, List<P> params, int groupSize, int timeout, Executor executor) {
        if (CollectionUtils.isEmpty(params)) {
            return Collections.emptyList();
        }
        if (groupSize <= 0) {
            groupSize = DEFAULT_SIZE;
        }
        if (timeout <= 0) {
            timeout = DEFAULT_TIMEOUT;
        }
        if (Objects.isNull(executor)) {
            throw new RuntimeException("批量执行方法(" + function.getClass().getSimpleName() + ")异常,executor is null");
        }
        // 如果小于指定批量则直接使用主线程进行操作
        if (params.size() <= groupSize) {
            return syncExecute(function, params);
        }
        return asyncExecute(function, params, groupSize, timeout, executor);
    }


    /**
     * 使用主线程同步执行
     */
    private static <P, R> List<R> syncExecute(Function<List<P>, R> function, List<P> params) {
        List<R> results = Lists.newArrayListWithCapacity(params.size());
        R r = function.apply(params);
        if (r != null) {
            results.add(r);
        }
        return results;
    }

    /**
     * 使用线程池异步执行
     */
    private static <P, R> List<R> asyncExecute(Function<List<P>, R> function, List<P> params, int groupSize, int timeout, Executor executor) {
        // 拆分批量多线程处理
        List<R> results = Collections.synchronizedList(Lists.newArrayListWithCapacity(params.size()));
        List<List<P>> partitionList = Lists.partition(params, groupSize);
        CountDownLatch latch = new CountDownLatch(partitionList.size());
        for (List<P> strings : partitionList) {
            doWithChildThread(function, strings, results, latch, executor);
        }
        try {
            boolean v = latch.await(timeout, TimeUnit.MILLISECONDS);
            if (!v) {
                log.warn("批量任务并发执行,请求总时间超过{}ms,返回null", timeout);
                return results;
            }
        } catch (InterruptedException e) {
            log.error("批量任务并发信息失败，参数信息={}", params, e);
            Thread.currentThread().interrupt();
        }
        return results;
    }


    //Child Thread
    private static <P, R> void doWithChildThread(Function<List<P>, R> function, List<P> params, List<R> lists, CountDownLatch latch, Executor executor) {
        executor.execute(() -> {
            try {
                R result = function.apply(params);
                if (Objects.nonNull(result)) {
                    lists.add(result);
                }
            } finally {
                latch.countDown();
            }
        });
    }
}
